# -*- coding: utf-8 -*-
# @Time    : 2020/6/21 下午9:14
# @Author  : caotian
# @FileName: residualnetmyopiamodel.py
# @Software: PyCharm

import numpy as np
import paddle
import paddle.fluid as fluid
from paddle.fluid.dygraph.nn import Linear,Conv2D,Pool2D,BatchNorm
from paddle.fluid.layer_helper import LayerHelper
from paddle.fluid.dygraph.base import to_variable

class ConvBNLayer(fluid.dygraph.Layer):
    def __init__(self,
                 num_channels,
                 num_filters,
                 filter_size,
                 stride=1,
                 groups=1,
                 act=None):
        super(ConvBNLayer,self).__init__()
        self._conv=Conv2D(num_channels=num_channels,
                          num_filters=num_filters,
                          filter_size=filter_size,
                          stride=stride,
                          padding=(filter_size-1)//2,
                          groups=groups,
                          act=None,
                          bias_attr=False)
        self._batch_norm=BatchNorm(num_filters,act=act)
    def forward(self,inputs):
        y=self._conv(inputs)
        y=self._batch_norm(y)
        return y
class BottleneckBlock(fluid.dygraph.Layer):
    def __init__(self,num_channels,num_filters,stride,shortcut=True):
        super(BottleneckBlock,self).__init__()
        self.conv0=ConvBNLayer(num_channels=num_channels,num_filters=num_filters,filter_size=1,act='relu')
        self.conv1=ConvBNLayer(num_channels=num_filters,num_filters=num_filters,filter_size=3,stride=stride,act='relu')
        self.conv2=ConvBNLayer(num_channels=num_filters,num_filters=num_filters*4,filter_size=1,act=None)
        if not shortcut:
            self.short=ConvBNLayer(num_channels=num_channels,num_filters=num_filters*4,filter_size=1,stride=stride)
        self.shortcut=shortcut
        self._num_channels_out=num_filters*4
    def forward(self,x):
        y=self.conv0(x)
        conv1=self.conv1(y)
        conv2=self.conv2(conv1)
        if self.shortcut:
            short=x
        else:
            short=self.short(x)
        y=fluid.layers.elementwise_add(x=short,y=conv2)
        layer_helper=LayerHelper(self.full_name(),act='relu')
        return layer_helper.append_activation(y)

class ResNet(fluid.dygraph.Layer):
    def __init__(self,layers=50,class_dim=1):
        super(ResNet, self).__init__()
        self.layers=layers
        supported_layers=[50,101,152]
        assert layers in supported_layers,"supported layers are {} but input layer is {}".format(supported_layers,layers)
        if layers==50:
            depth=[3,4,6,3]
        elif layers==101:
            depth=[3,4,23,3]
        else:
            depth=[3,8,36,3]
        num_filters=[64,128,256,512]
        self.conv=ConvBNLayer(num_channels=3,num_filters=64,filter_size=7,stride=2,act='relu')
        self.pool2d_max=Pool2D(pool_size=3,pool_stride=2,pool_padding=1,pool_type='max')
        self.bottleneck_block_list=[]
        num_channels=64
        for block in range(len(depth)):
            shortcut=False
            for i in range(depth[block]):
                bb=BottleneckBlock(num_channels=num_channels,num_filters=num_filters[block],stride=2 if i==0 and block != 0 else 1,shortcut=shortcut)
                bottleneckblock=self.add_sublayer('bb_%d_%d'%(block,i),bb)
                num_channels=bottleneckblock._num_channels_out
                self.bottleneck_block_list.append(bottleneckblock)
                shortcut=True
        self.pool2d_avg=Pool2D(pool_size=7,pool_type='avg',global_pooling=True)
        import  math
        stdv=1.0/math.sqrt(2048*1.0)
        self.out=Linear(input_dim=2048,output_dim=class_dim,param_attr=fluid.param_attr.ParamAttr(initializer=fluid.initializer.Uniform(-stdv,stdv)))
    def forward(self,inputs):
        y=self.conv(inputs)
        y=self.pool2d_max(y)
        for bottleneck_block in self.bottleneck_block_list:
            y=bottleneck_block(y)
        y=self.pool2d_avg(y)
        y=fluid.layers.reshape(y,[y.shape[0],-1])
        y=self.out(y)
        return y